package com.cloud.gateway.limiter;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
public class LimiterManager {

    private List<LimiterDefinition> limiterDefinitionList;
    private Map<String, Bucket> bucketMap;

    /**
     * 初始化
     * @param limiterDefinitionList
     */
    public void init(List<LimiterDefinition> limiterDefinitionList) {
        if (CollectionUtils.isEmpty(limiterDefinitionList)) {
            return;
        }
        Map<String, Bucket> bucketMap = new HashMap<>();
        for (LimiterDefinition limiterDefinition : limiterDefinitionList) {
            Refill refill = Refill.greedy(limiterDefinition.getRefillTokens(), Duration.ofMillis(limiterDefinition.getRefillPeriod()));
            Bandwidth limit = Bandwidth.classic(limiterDefinition.getCapacity(), refill);
            Bucket bucket = Bucket4j.builder().addLimit(limit).build();

            String key = limiterDefinition.getPath() + limiterDefinition.getLimitUserId();
            bucketMap.put(key, bucket);
        }
        this.limiterDefinitionList = limiterDefinitionList;
        this.bucketMap = bucketMap;
    }

    /**
     * 是否允许
     * @param requestURI
     * @param userId
     * @return boolean
     */
    public boolean isAllow(String requestURI, String userId) {
        if (!CollectionUtils.isEmpty(limiterDefinitionList)) {
            for (LimiterDefinition limiterDefinition : limiterDefinitionList) {
                if (requestURI.startsWith(limiterDefinition.getPath())) {
                    Bucket bucket = bucketMap.get(limiterDefinition.getPath() + userId);
                    if (bucket == null) {
                        bucket = bucketMap.get(limiterDefinition.getPath());
                    }
                    if (bucket == null) {
                        continue;
                    }
                    return bucket.tryConsume(1);
                }
            }
        }
        return true;
    }

}
